# vim: syntax=py
# -*- mode: python -*-
# SPDX-License-Identifier: Apache-2.0
# Copyright(c) 2019 Intel Corporation

# for errnos
import errno
from conf.utils import *


# ====================================================
#       Parameters
# ====================================================


conf_file = get_env('CONF_FILE', 'conf/spgwu.json')
conf = get_json_conf(conf_file, False)

# Maximum number of flows to manage ip4 frags for re-assembly
max_ip_defrag_flows = None
try:
    max_ip_defrag_flows = int(conf["max_ip_defrag_flows"])
except ValueError:
    print('Invalid value for max_ip_defrag_flows. Not installing IP4Defrag module.')
except KeyError:
    print('max_ip_defrag_flows value not set. Not installing IP4Defrag module.')

# Enable ip4 fragmentation
ip_frag_with_eth_mtu = None
try:
    ip_frag_with_eth_mtu = int(conf["ip_frag_with_eth_mtu"])
except ValueError:
    print('Invalid value for ip_frag_with_eth_mtu. Not installing IP4Frag module.')
except KeyError:
    print('ip_frag_with_eth_mtu value not set. Not installing IP4Frag module.')

# See this link for details:
## https://github.com/NetSys/bess/blob/master/bessctl/module_tests/timestamp.py
measure = False
try:
    measure = bool(conf["measure"])
except ValueError:
    print('Invalid value for measure. Not installing Measure module.')
except KeyError:
    print('measure value not set. Not installing Measure module.')

# Detect mode. Default is dpdk
mode = None
try:
    mode = conf["mode"]
except KeyError:
    print('Autodetecting network driver')


enable_ntf = False
try:
    enable_ntf = bool(conf["enable_ntf"])
except ValueError:
    print('Invalid value for enable_ntf. Disabling Network Token Functions.')
except KeyError:
    print('enable_ntf value not set. Disabling Network Token Functions.')

# ====================================================
#       Port Helpers
# ====================================================


MAX_GATES = 8192
dpdk_ports = {}


def scan_dpdk_ports():
    idx = 0
    while True:
        try:
            intf = PMDPort(name="Port {}".format(idx), port_id=idx)
            if intf:
                # Need to declare mac so that we don't lose key during destroy_port
                mac = intf.mac_addr
                dpdk_ports[mac] = idx
                bess.destroy_port(intf.name)
        except bess.Error as e:
            if e.code == errno.ENODEV:
                break
            else:
                raise
        idx += 1
        # RTE_MAX_ETHPORTS is 32 and we need 2 for vdevs
        if idx == 30:
          break
    return True if dpdk_ports else False


class Port:
    def __init__(self, name, ext_addrs):
        self.name = name
        self.wid = None
        self.fpi = None
        self.fpo = None
        self.bpf = None
        self.bpfgate = 0
        self.routes_table = None
        self.nat = None
        self.ext_addrs = ext_addrs

    def bpf_gate(self):
        if self.bpfgate < MAX_GATES - 2:
            self.bpfgate += 1
            return self.bpfgate
        else:
            raise Exception('Port {}: Out of BPF gates to allocate'.format(self.name))

    def detect_mode(self):
        try:
            peer_by_interface(self.name)
            mode = 'dpdk'
        except:
            mode = 'linux'
        return mode

    def init_fastpath(self, **kwargs):
        # Initialize PMDPort and RX/TX modules
        name = self.name
        fast = PMDPort(name="{}Fast".format(name), **kwargs)
        self.fpi = __bess_module__("{}FastPI".format(name), 'PortInc', port=fast.name)
        self.fpo = __bess_module__("{}FastPO".format(name), 'PortOut', port=fast.name)

        # Initialize BPF to classify incoming traffic to go to kernel and/or pipeline
        self.bpf = __bess_module__("{}FastBPF".format(name), 'BPF')
        self.bpf.clear()

        # Default drop when no matches
        if measure:
            self.fpi -> Timestamp() -> self.bpf:0 -> Sink()
        else:
            self.fpi -> self.bpf:0 -> Sink()

        # Initialize route module
        self.rtr = __bess_module__("{}Routes".format(name), 'IPLookup')

        # Default route goes to Sink
        self.rtr.add(prefix='0.0.0.0', prefix_len=0, gate=MAX_GATES-1)
        self.rtr:(MAX_GATES-1) -> Sink()

        # Attach fastpath to worker's root TC
        self.fpi.attach_task(wid=self.wid)

    def setup_port(self, idx, workers):
        # Pick the worker handling this port
        self.wid = idx % workers

        name = self.name
        wid = self.wid
        print('Setting up port {} on worker {}'.format(name,wid))

        # Detect the mode of this interface - DPDK/AF_XDP/AF_PACKET
        global mode
        if mode is None:
            mode = self.detect_mode()

        if mode not in ['af_xdp', 'linux', 'dpdk', 'af_packet']:
            raise Exception('Invalid mode: {} selected.'.format(mode))

        if mode in ['af_xdp', 'linux']:
            try:
                # Initialize kernel fastpath.
                # AF_XDP requires that num_rx_qs == num_tx_qs
                kwargs = {"vdev" : "net_af_xdp{},iface={},start_queue=0,queue_count={}"
                          .format(idx, name, workers), "num_out_q": workers, "num_inc_q": workers}
                self.init_fastpath(**kwargs)
            except:
                if mode == 'linux':
                    print('Failed to create AF_XDP socket for {}. Retrying with AF_PACKET socket...'.format(name))
                    mode = 'af_packet'
                else:
                    print('Failed to create AF_XDP socket for {}. Exiting...'.format(name))
                    sys.exit()

        if mode == 'af_packet':
            try:
                # Initialize kernel fastpath
                kwargs = {"vdev" : "net_af_packet{},iface={},qpairs={}".format(idx, name, workers), "num_out_q": workers}
                self.init_fastpath(**kwargs)
            except:
                print('Failed to create AF_PACKET socket for {}. Exiting...'.format(name))
                sys.exit()

        if mode == 'dpdk':
            kwargs = None
            pci = alias_by_interface(name)
            if pci is not None:
                kwargs = {"pci": pci, "num_out_q": workers}
                try:
                    self.init_fastpath(**kwargs)
                except:
                    kwargs = None
                    print('Unable to initialize {} fastpath using alias {},\
                        falling back to scan'.format(name, pci))
            if kwargs is None:
                # Fallback to scanning ports
                # if port list is empty, scan for dpdk_ports first
                if not dpdk_ports and scan_dpdk_ports() == False:
                    print('Registered dpdk ports do not exist.')
                    sys.exit()
                # Initialize DPDK fastpath
                fidx = dpdk_ports.get(mac_by_interface(name))
                if fidx is None:
                    raise Exception(
                        'Registered port for {} not detected!'.format(name))
                kwargs = {"port_id": fidx, "num_out_q": workers}
                self.init_fastpath(**kwargs)

            # Initialize kernel slowpath port and RX/TX modules
            try:
                peer = peer_by_interface(name)
                vdev = "net_af_packet{},iface={}".format(idx, peer)
                slow = PMDPort(name="{}Slow".format(name), vdev=vdev)
                spi = __bess_module__("{}SlowPI".format(name), 'PortInc', port=slow.name)
                spo = __bess_module__("{}SlowPO".format(name), 'PortOut', port=slow.name)
                qspo = __bess_module__("{}QSlowPO".format(name), 'Queue')

                # host_ip_filter: tcpdump -i foo 'dst host 198.19.0.1 or 198.18.0.1' -d
                # Should always be set to lowest priority
                HostGate = MAX_GATES - 1
                ips = ips_by_interface(name)
                host_ip_filter = {"priority": -HostGate, "filter": "dst host "
                                + " or ".join(str(x) for x in ips), "gate": HostGate}

                self.bpf.add(filters=[host_ip_filter])

                # Direct control traffic from DPDK to kernel
                self.bpf:HostGate -> qspo -> spo

                # Direct control traffic from kernel to DPDK
                spi -> self.fpo

                tc = 'slow{}'.format(wid)
                try:
                    bess.add_tc(tc, policy='round_robin', wid=wid)
                except Exception as e:
                    if e.errmsg == "Name '{}' already exists".format(tc):
                        pass
                    else:
                        raise e
                # Limit scheduling slow path RX/TX to 1000 times/second each
                for mod in spi, qspo:
                    bess.add_tc(mod.name,
                            parent=tc,
                            policy='rate_limit',
                            resource='count',
                            limit={'count': 1000})
                    mod.attach_task(mod.name)
            except Exception as e:
                print('Mirror veth interface: {} misconfigured: {}'.format(name, e))

        # Direct fast path traffic to Frag module
        merge = __bess_module__("{}Merge".format(name), 'Merge')

        out = self.fpo

        # Attach frag module (if enabled)
        if ip_frag_with_eth_mtu is not None:
            frag = __bess_module__("{}IP4Frag".format(name), 'IPFrag', mtu=ip_frag_with_eth_mtu)
            frag:1 -> out
            frag:0 -> Sink()
            out = frag

        # Attach telemeric module (if enabled)
        if measure:
            m = Measure()
            m -> out
            out = m

        # Attach nat module (if enabled)
        if self.ext_addrs is not None:
            # Tokenize the string
            addrs = self.ext_addrs.split(' or ')
            # Make a list of ext_addr
            nat_list = list()
            for addr in addrs:
                nat_dict = dict()
                nat_dict['ext_addr'] = addr
                nat_list.append(nat_dict)

            # Create the NAT module
            self.nat = __bess_module__("{}NAT".format(conf[interface]["ifname"]), 'NAT', ext_addrs=nat_list)
            self.nat:1 -> out
            out = self.nat

        # Attach Merge module to the 'outlist' of modules
        merge -> out


# ====================================================
#       Validate
# ====================================================


# CIDRs
enb_cidr = conf["enb_cidr"]
ue_cidr = conf["ue_cidr"]

cidrs = [enb_cidr, ue_cidr]
for cidr in cidrs:
    if validate_cidr(cidr) is False:
        exit(1, 'Invalid {}'.format(cidr))

# ====================================================
#       Core Setup
# ====================================================


# Initialize workers
workers = 1
try:
    workers = int(conf["workers"])
except ValueError:
    print('Invalid workers value! Re-setting # of workers to 1.')

cores = get_process_affinity()
for wid in xrange(workers):
    bess.add_worker(wid=wid, core=int(cores[wid % len(cores)]))

# ====================================================
#       Port Setup
# ====================================================


interfaces = ["s1u", "sgi"]
ports = {}

for idx, interface in enumerate(interfaces):
    try:
        ext_addrs = conf[interface]["ip_masquerade"]
    except KeyError:
        ext_addrs = None

    port = Port(conf[interface]["ifname"], ext_addrs)
    if port.name in ports:
        continue
    port.setup_port(idx, workers)
    ports[port.name] = port

s1u_ifname = conf["s1u"]["ifname"]
sgi_ifname = conf["sgi"]["ifname"]


# ====================================================
#       Network Token Functions Setup
# ====================================================

ntf = None
if enable_ntf:
    ntf = NTF()
    # TODO: Populate these tables by talking to the NTE
    ntf.table_create(dpid=1, max_entries=5)
    ntf.entry_create(dpid=1, entry_id=1, dscp=1, token={
      'app_id': 0x10001,
      'encryption_key': '{"alg":"A128CBC-HS256","k":"Qr_XwDGctna3SlR88rEJYt6Zm100SASYeJWSJihDnsA","key_ops":["encrypt","decrypt"],"kty":"oct"}'
    })


# ====================================================
#       Downlink Pipeline
# ====================================================

s1uRoutes = ports[s1u_ifname].rtr

# Maximum number of sessions to manage
try:
    max_sessions = int(conf["max_sessions"])
except ValueError:
    print('Invalid max_sessions value!')

# Setting filter to detect UE subnet
sgiFastBPF = ports[sgi_ifname].bpf
UEGate = ports[sgi_ifname].bpf_gate()
if ports[sgi_ifname].nat is not None:
    ue_filter = {"priority": -UEGate,
                 "filter": "ip dst {}".format(ports[sgi_ifname].ext_addrs), "gate": UEGate}
else:
    ue_filter = {"priority": -UEGate,
                 "filter": "ip dst net {}".format(ue_cidr), "gate": UEGate}
sgiFastBPF.add(filters=[ue_filter])

# Verify IP/UDP Checksum
sgiFastBPF:UEGate -> sgiRxIPCksum::IPChecksum(verify=True) -> sgiRxUDPCksum::L4Checksum(verify=True)
_in = sgiRxUDPCksum
gate = 0
sgiRxIPCksum:1 -> Sink()
sgiRxUDPCksum:1 -> Sink()

# Append nat module (if enabled)
if ports[sgi_ifname].nat is not None:
    _in:gate -> 1:ports[sgi_ifname].nat
    _in = ports[sgi_ifname].nat
    gate = 0

# Append sgiIP4Defrag module (if enabled)
if max_ip_defrag_flows is not None:
    _in:gate -> sgiIP4Defrag::IPDefrag(num_flows=max_ip_defrag_flows, numa=0)
    _in = sgiIP4Defrag
    gate = 1
    # Drop pkts that are unable to be fragmented
    sgiIP4Defrag:0 -> Sink()

_in:gate -> EtherTrim::GenericDecap(bytes=14)

GTPUEncap::GtpuEncap(s1u_sgw_ip=ip2long(ips_by_interface(s1u_ifname)[0]), num_subscribers=max_sessions):1 \
    -> S1UEtherAdd::GenericEncap(fields=[
        {'size': 6, 'value': {'value_int': 0x0}},
        {'size': 6, 'value': {'value_int': mac2hex(mac_by_interface(s1u_ifname))}},
        {'size': 2, 'value': {'value_int': 0x0800}}]) \
    -> OuterUDPCsum::L4Checksum() \
    -> OuterIPCsum::IPChecksum() \
    -> s1uRoutes

_in = EtherTrim
if enable_ntf:
    _in -> 1:ntf
    _in = ntf
_in:0 -> GTPUEncap

# Drop unknown packets
GTPUEncap:0 -> Sink()
OuterIPCsum:1 -> Sink()
OuterUDPCsum:1 -> Sink()



# ====================================================
#       Uplink Pipeline
# ====================================================


# Setting filter to detect gtpu traffic
# src net 11.1.1.0 mask 255.255.255.0           # check eNB subnet
# and dst host 11.1.1.1                         # check S/PGWU IP
# and udp dst port 2152                         # check GTPU port
# and (udp[28:4] & 0xffffff00) = 0x10000000     # check UE subnet
s1uFastBPF = ports[s1u_ifname].bpf
check_enb_subnet = "src net {} ".format(enb_cidr)
check_spgwu_ip = " and dst host " + \
    " or ".join(str(x) for x in ips_by_interface(s1u_ifname))
check_gtpu_port = " and udp dst port 2152"
check_ue_subnet = " and (udp[28:4] & 0x{}) = 0x{}".format(
    cidr2mask(ue_cidr), ip2hex(cidr2block(ue_cidr)[0]))
check_gtpu_msg_echo = " and udp[9] = 0x1"

GTPUEchoGate = ports[s1u_ifname].bpf_gate()
uplink_echo_filter = {"priority": -GTPUEchoGate, "filter": check_enb_subnet +
                      check_spgwu_ip + check_gtpu_port +
                      check_gtpu_msg_echo, "gate": GTPUEchoGate}
s1uFastBPF.add(filters=[uplink_echo_filter])

GTPUGate = ports[s1u_ifname].bpf_gate()
uplink_filter = {"priority": -GTPUGate, "filter": check_enb_subnet +
               check_spgwu_ip + check_gtpu_port, "gate": GTPUGate}
s1uFastBPF.add(filters=[uplink_filter])

sgiRoutes = ports[sgi_ifname].rtr

# Verify IP Checksum
s1uFastBPF:GTPUGate -> s1uRxIPCksum::IPChecksum(verify=True) -> s1uRxUDPCksum::L4Checksum(verify=True)
s1uRxIPCksum:1 -> Sink()
s1uRxUDPCksum:1 -> Sink()

# Append s1uIP4Defrag module (if enabled)
if max_ip_defrag_flows is not None:
    s1uRxUDPCksum:0 -> s1uIP4Defrag::IPDefrag(num_flows=max_ip_defrag_flows, numa=0):1 -> EtherDecapTrim::GenericDecap(bytes=14)
    # Drop pkts that are unable to be fragmented
    s1uIP4Defrag:0 -> Sink()
else:
    s1uRxUDPCksum:0 -> EtherDecapTrim::GenericDecap(bytes=14)

EtherDecapTrim -> GTPUDecap::GtpuDecap(ename="GTPUEncap")

SGIEtherAdd::GenericEncap(fields=[
        {'size': 6, 'value': {'value_int': 0x0}},
        {'size': 6, 'value': {'value_int': mac2hex(mac_by_interface(sgi_ifname))}},
        {'size': 2, 'value': {'value_int': 0x0800}}]) \
    -> sgiRoutes

_in = GTPUDecap
if enable_ntf:
    _in:1 -> 0:ntf
    _in = ntf
_in:1 -> SGIEtherAdd

s1uFastBPF:GTPUEchoGate \
        -> gtpRxIPCksum::IPChecksum(verify=True) \
        -> gtpRxUDPCksum::L4Checksum(verify=True) \
        -> GTPUEcho::GtpuEcho(s1u_sgw_ip=ip2long(ips_by_interface(s1u_ifname)[0])):1 \
        -> EthSwap::MACSwap() \
        -> 1:OuterUDPCsum

# Drop unknown packets
GTPUEcho:0 -> Sink()
GTPUDecap:0 -> Sink()
gtpRxIPCksum:1 -> Sink()
gtpRxUDPCksum:1 -> Sink()


# ====================================================
#       SIM_TEST
# ====================================================

import time

def sim_start_test():
    start_teid = 0xf0000000
    start_ue_ip = 0x10000001
    start_enb_ip = 0x0b010181
    NG4T_MAX_UE_RAN = 500000
    NG4T_MAX_ENB_RAN = 80

    for i in range(max_sessions):

        # NG4T-based formula to calculate enodeB IP address against a given UE IP address
        # il_trafficgen also uses the same scheme
        # See SimuCPEnbv4Teid(...) in ngic code for more details
        ue_of_ran = i % NG4T_MAX_UE_RAN
        ran = i / NG4T_MAX_UE_RAN
        enb_of_ran = ue_of_ran % NG4T_MAX_ENB_RAN
        enb_idx = ran * NG4T_MAX_ENB_RAN + enb_of_ran

        GTPUEncap.add(teid=start_teid+i, eteid=i+1, ueaddr=start_ue_ip+i, enodeb_ip=start_enb_ip+enb_idx)

def sim_end_test():
    start_ue_ip = 0x10000001

    for i in range(max_sessions):
        GTPUEncap.remove(ueaddr=start_ue_ip+i)

    GTPUEncap.show_records()

### Uncomment the following lines to test with il_trafficgen ###
# Start the test
#sim_start_test()
# Wait for 30 seconds before deleting the session info entries (optional)
#time.sleep(30)
#sim_end_test()

# Finally send SIGHUP to route_control daemon on reload
# TODO: behavior is unspecified if route_control.py pid is not found
route_control_pid = getpythonpid('route_control.py')
if route_control_pid:
    os.kill(route_control_pid, signal.SIGHUP)
